from mmengine.config import Config
from mmengine.analysis import get_model_complexity_info
from mmdet.registry import MODELS
import torch

# 动态导入自定义模块
from mmengine.utils import import_modules_from_strings

custom_imports = [
    'projects.DiffusionDet.diffusiondet',
]
import_modules_from_strings(custom_imports)

# 加载配置文件
config_file = '/media/ross/8TB/project/lsh/deep_learning/DiffusionDet_mmdet/DiffusionDet/work_dirs/ablation/config/diffusiondet_r50_lamfpn8_epoch_microalgeaOri_1lcm2_1adem2_1ddim4_1distill4_memeryOptim.py'
cfg = Config.fromfile(config_file)

# 构建模型
model = MODELS.build(cfg.model)
model.eval()

# 设置输入形状
input_shape = (3, 800, 1216)  # 通道数、高度、宽度

# 自定义 forward 函数
def custom_forward(model, inputs):
    init_bboxes = torch.zeros((1, 100, 4)).to(inputs.device)  # 根据需求调整
    init_t = torch.zeros((1, 100)).to(inputs.device)          # 根据需求调整
    return model.bbox_head.forward(inputs, init_bboxes, init_t)

# 统计 FLOPs 和参数量
flops, params = get_model_complexity_info(
    model, input_shape, as_strings=True, print_per_layer_stat=False, custom_forward=custom_forward
)
print(f"FLOPs: {flops}")
print(f"Params: {params}")
